﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using TracedReality.Primitive;

namespace TracedReality.RayTracing.PhotonMapping
{
    /// <summary>
    /// A basic kd-tree implementation for holding Photons.
    /// </summary>
    public class PhotonTree
    {
        private BoundingBox bound;
        private PhotonTree left;
        private PhotonTree right;
        private int depth;
        private List<Photon> members;

        /// <summary>
        /// Create a node in the tree with the given BoundingBox and depth.
        /// </summary>
        /// <param name="box">The BoundingBox for this node.</param>
        /// <param name="d">The depth of this node.</param>
        public PhotonTree(BoundingBox box, int d = 0)
        {
            bound = box;
            depth = d;
            left = null;
            right = null;
            members = new List<Photon>();
        }

        /// <summary>
        /// Check if this node's BoundingBox overlaps another BoundingBox.
        /// </summary>
        /// <param name="b">The oher BoundingBox to check.</param>
        /// <returns>True if they overlap, false otherwise.</returns>
        public bool overlaps(BoundingBox b) { return bound.overlaps(b); }

        /// <summary>
        /// Check if this node is a leaf node.
        /// </summary>
        /// <returns>True if this node is a leaf, false otherwise.</returns>
        public bool isLeaf() { return (left == null && right == null); }

        /// <summary>
        /// Get the height of the tree represented by this node.
        /// </summary>
        /// <returns>The height of the tree under and including this node.</returns>
        public int getHeight()
        {
            if (isLeaf()) { return 1; }

            return 1 + Math.Max(left.getHeight(), right.getHeight());
        }

        /// <summary>
        /// Get a count of how many Photons are represented by this tree.
        /// </summary>
        /// <returns>How many Photons are stored in this tree.</returns>
        public int getCount()
        {
            if (isLeaf()) { return members.Count; }

            return left.getCount() + right.getCount();
        }

        /// <summary>
        /// Get all the Photons that are in BoundingBoxes that overlap the given BoundingBox.
        /// </summary>
        /// <param name="b">The BoundingBox to find Photons in.</param>
        /// <returns>A List of Photons that may be in the given BoundingBox.</returns>
        public List<Photon> getInBox(BoundingBox b)
        {
            List<Photon> result = new List<Photon>();
            List<PhotonTree> todo = new List<PhotonTree>();
            todo.Add(this);
            do
            {
                PhotonTree cur = todo[todo.Count - 1];
                todo.RemoveAt(todo.Count - 1);

                if (!cur.overlaps(b)) { continue; }
                if (!cur.isLeaf())
                {
                    todo.Add(cur.left);
                    todo.Add(cur.right);
                }
                else
                {
                    result.AddRange(cur.members);
                }
            } while (todo.Count != 0);

            return result;
        }

        /// <summary>
        /// Add a Photon to this tree.
        /// </summary>
        /// <param name="item">Photon to add.</param>
        public void add(Photon item)
        {
            Vector p = item.position;
            if (isLeaf())
            {
                if (bound.contains(p))
                {
                    members.Add(item);
                }

                return;
            }

            left.add(item);
            right.add(item);
        }

        /// <summary>
        /// Used for sorting the Photons along a certain axis.
        /// </summary>
        private class PhotonComparer : IComparer<Photon>
        {
            public int axis;

            public int Compare(Photon a, Photon b)
            {
                int ret = 0;
                switch (axis)
                {
                    case 0 :
                        ret = a.position.x.CompareTo(b.position.x);
                        break;
                    case 1 :
                        ret = a.position.y.CompareTo(b.position.y);
                        break;
                    case 2 :
                        ret = a.position.z.CompareTo(b.position.z);
                        break;
                }

                return ret;
            }
        }

        /// <summary>
        /// Find the value along the given axis at which to split the node.
        /// This uses the median of the data along the axis to ensure a balanced tree.
        /// </summary>
        /// <param name="axis">Which axis to split along.</param>
        /// <returns>The value to split at.</returns>
        private double selectSplitValue(int axis, Random r)
        {
            PhotonComparer pc = new PhotonComparer();
            pc.axis = axis;

            //members.Sort(pc);

            double val = 0;
            Photon mid = findKMedian(members, members.Count / 2, pc, r);
            //mid = members[members.Count / 2];
            switch (axis)
            {
                case 0 :
                    val = mid.position.x;
                    break;
                case 1 :
                    val = mid.position.y;
                    break;
                case 2 :
                    val = mid.position.z;
                    break;
            }

            return val;
        }

        /// <summary>
        /// Finds the k-th median of a given set of Photons. For the overall median,
        /// k = A.Count / 2.
        /// </summary>
        /// <param name="A">Set of Photons.</param>
        /// <param name="K">Which median to find.</param>
        /// <param name="pc">Comparator.</param>
        /// <param name="r">Random for pivot choice.</param>
        /// <returns>The k-th median</returns>
        private Photon findKMedian(List<Photon> A, int K, PhotonComparer pc, Random r)
        {
            Photon pivot = A[r.Next(A.Count)];

            List<Photon> S = new List<Photon>();
            List<Photon> B = new List<Photon>();
            foreach (Photon p in A)
            {
                int cmp = pc.Compare(p, pivot);
                if (cmp < 0) { S.Add(p); }
                else if (cmp > 0) { B.Add(p); }
            }

            if (S.Count == (K - 1) || (S.Count == 0 && B.Count == 0)) { return pivot; }
            if (S.Count < (K - 1) && B.Count > 0) { return findKMedian(B, K - S.Count - 1, pc, r); }
            else if (S.Count > 0) { return findKMedian(S, K, pc, r); }
            else { throw new Exception(); }
        }

        /// <summary>
        /// Split the current node if it is a leaf so that is contains no more than
        /// a given amount as long as it is above a given depth.
        /// </summary>
        /// <param name="maxcap">The maximum capacity of a node.</param>
        /// <param name="maxdep">The maximum depth of a node.</param>
        /// <param name="r">A random used when finding the median.</param>
        public void splitCell(int maxcap, int maxdep, Random r)
        {
            if (!isLeaf()) { return; }
            if (depth >= maxdep) { return; }
            if (members.Count < maxcap) { return; }

            Vector max = bound.max;
            Vector min = bound.min;
            double dx = max.x - min.y;
            double dy = max.y - min.y;
            double dz = max.z - min.z;

            Vector min1 = new Vector(min);
            Vector min2 = new Vector(min);
            Vector max1 = new Vector(max);
            Vector max2 = new Vector(max);

            int splitAxis;
            if (dx >= dy && dx >= dz)
            {
                splitAxis = 0;
            }
            else if (dy >= dx && dy >= dz)
            {
                splitAxis = 1;
            }
            else
            {
                splitAxis = 2;
            }

            double splitValue = selectSplitValue(splitAxis, r);

            switch (splitAxis)
            {
                case 0:
                    max1.x = min2.x = splitValue;
                    break;
                case 1:
                    max1.y = min2.y = splitValue;
                    break;
                case 2:
                    max1.z = min2.z = splitValue;
                    break;
            }

            left = new PhotonTree(new BoundingBox(min1, max1), depth + 1);
            right = new PhotonTree(new BoundingBox(min2, max2), depth + 1);

            foreach (Photon p in members) { add(p); }
            members = null;

            left.splitCell(maxcap, maxdep, r);
            right.splitCell(maxcap, maxdep, r);
        }
    }
}
